import pandas as pd
import matplotlib.pyplot as plt
from scipy import stats
import seaborn as sns
import numpy as np
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
import argparse


def plot_corr_heatmap(dataset, target, sort_num=10):
    corr = dataset.corr()
    corr_with_target = corr[target].sort_values(ascending=False)
    top_ten_factors = corr_with_target[:sort_num]
    top_ten_matrix = dataset[top_ten_factors.index].corr()
    plt.figure(figsize=(20, 16))
    sns.heatmap(top_ten_matrix, annot=True, cmap='coolwarm')
    plt.savefig('plot.png')

def pearsonr_analysis(dataset, target):
    for i in range(dataset.shape[1]):
        if dataset.columns[i] == target:
            continue
        cor = stats.pearsonr(dataset.loc[:,target],dataset.iloc[:,i])
        print("%s pearsoner: %5f %5f"%(dataset.columns[i].ljust(20), cor.statistic, cor.pvalue))
        if abs(cor.statistic) >= 0.3:
            plt.figure(dataset.columns[i]+'-'+target,figsize=(20,20), dpi=70)
            plt.subplot(2,2,1)
            plt.scatter(dataset.iloc[:,i],dataset.loc[:,target])
            plt.xlabel(dataset.columns[i], fontdict={'size': 25})
            plt.ylabel(target, fontdict={'size': 25})
            plt.subplot(2,2,2)
            plt.hist(dataset.iloc[:,i],bins='auto')
            plt.xlabel(dataset.columns[i], fontdict={'size': 25})
            plt.subplot(2,1,2)
            x=range(0,dataset.shape[0])
            plt.plot(x,dataset.iloc[:,i],label=dataset.columns[i],linewidth=1,color='r')
            plt.plot(x,dataset.loc[:,target],label=target,linewidth=1)
            plt.legend()
            plt.savefig('plot.png')

def pca_scree_plot(dataset):
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(dataset)
    pca = PCA()
    pca.fit(X_scaled)
    explained_variance = pca.explained_variance_ratio_
    cumulative_variance = np.cumsum(explained_variance)
    components_count = np.argmax(cumulative_variance>=0.95)+1

    # 绘制屏幕图
    plt.figure(figsize=(10, 5))
    plt.bar(range(1, len(explained_variance) + 1), explained_variance, alpha=0.5, align='center', label='Individual explained variance')
    plt.step(range(1, len(explained_variance) + 1), np.cumsum(explained_variance), where='mid', label='Cumulative explained variance')
    plt.axhline(y=0.95, color='r', linestyle='-', label='95% explained variance')
    plt.axvline(x=components_count, color='b', linestyle='--', label='Components for 95% variance')
    plt.text(components_count, 0.95, f' {components_count}components', color = 'blue', fontsize=12)
    plt.ylabel('Explained Variance Ratio')
    plt.xlabel('Individual Component Index')
    plt.title('PCA Scree Plot')
    plt.legend(loc='best')
    plt.tight_layout()
    plt.savefig('plot.png')

def target_diff(dataset, target):
    dataset['abs_diff'] = dataset[target].diff().abs()
    counts, bin_edges = np.histogram(dataset['abs_diff'].dropna(), bins=80)
    cumulative_counts = np.cumsum(counts)
    total_counts = cumulative_counts[-1]
    plt.hist(dataset['abs_diff'].dropna(), bins=80, edgecolor='black')
    for i in range(len(bin_edges) - 1):
        if i==0:continue
        if bin_edges[i] % 0.1 < 0.01: 
            plt.annotate(f'{cumulative_counts[i]/total_counts*100:.2f}%', 
                        xy=(bin_edges[i]-1/600, counts[i-1]), 
                        xytext=(0, 5),
                        textcoords='offset points',
                        ha='center', va='bottom')
    plt.title('The difference between adjacent rows in the target column')
    plt.savefig('plot.png')

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--input', default="../../data/prediction_data.csv", help='file to load')
    parser.add_argument('-t', '--target_feature', default="bucket_deg", help='target feature')
    parser.add_argument('-p', '--option', default="heatmap", help='analysis option')

    dataset = df = pd.read_table(parser.parse_args().input, sep=',')
    target = parser.parse_args().target_feature
    option = parser.parse_args().option

    if option == 'heatmap':
        plot_corr_heatmap(dataset, target)
    elif option == 'pearsonr':
        pearsonr_analysis(dataset, target)
    elif option == 'pca':
        pca_scree_plot(dataset)
    elif option == 'diff':
        target_diff(dataset, target)